package top.lileix.mybatis.generator.x.db;

import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import javax.sql.DataSource;

import top.lileix.mybatis.generator.x.domain.Column;
import top.lileix.mybatis.generator.x.domain.Table;
import top.lileix.mybatis.generator.x.kit.NamingKit;

public class MetaBuilder {
    protected DataSource dataSource;
    protected List<String> prefix;

    public MetaBuilder(DataSource dataSource, List<String> prefix) {
	if (dataSource == null) {
	    throw new IllegalArgumentException("dataSource can not be null.");
	}
	this.dataSource = dataSource;
	this.prefix = prefix;
    }

    /**
     * 获取所有的表
     * 
     * @throws SQLException
     */
    public List<Table> getAllTable() throws SQLException {
	List<Table> tables = new ArrayList<Table>();
	Connection conn = dataSource.getConnection();
	ResultSet rs = conn.getMetaData().getTables(null, null, "", null);
	String tableName = null;
	String remark = null;
	while (rs.next()) {
	    tableName = rs.getString("TABLE_NAME");
	    remark = rs.getString("REMARKS");
	    Table table = new Table();
	    table.setTableName(tableName);
	    table.setRemark(remark);
	    table.setModeName(NamingKit.toCapitalizeCamelCase(trim(tableName)));
	    getColumns(table, conn);
	    tables.add(table);
	}
	return tables;

    }

    /**
     * 去掉前缀
     * 
     * @param str
     * @return
     */
    public String trim(String str) {
	if (prefix != null) {
	    for (String string : prefix) {
		if (str.contains(string)) {
		    return str.replaceFirst(string, "");
		}
	    }
	}
	return str;
    }

    public void getColumns(Table table, Connection conn) throws SQLException {
	DatabaseMetaData dmd = conn.getMetaData();
	ResultSet pk = dmd.getPrimaryKeys(conn.getCatalog(), null, table.getTableName());
	String pkName = null;
	while (pk.next()) {
	    pkName = pk.getString("COLUMN_NAME");

	}
	ResultSet rs = dmd.getColumns(conn.getCatalog(), "%", table.getTableName(), "%");
	Set<Column> columns = new HashSet<Column>();
	while (rs.next()) {
	    Column c = new Column();
	    String column = rs.getString("COLUMN_NAME");
	    if (column.equals(pkName)) {
		c.setPrimaryKey(true);
		c.setRemark("主键");
	    }
	    c.setColumnName(column);
	    c.setFieldName(NamingKit.toCamelCase(column));
	    String columnType = rs.getString("TYPE_NAME");
	    c.setColumnType(columnType);
	    c.setFieldType(TypeMapping.getType(columnType));
	    String remark = rs.getString("REMARKS");

	    c.setRemark((c.getRemark() == null ? "" : c.getRemark()) + remark);

	    columns.add(c);
	}
	table.setColumns(columns);

    }
}
